#include <stdlib.h>
#include <time.h>
#include <unistd.h>
#include <fcntl.h>
#include <stdio.h>
#include <string.h>
#include <sys/time.h>

#include <sm2.h>
#include <sm3.h>
#include <sm4.h>

#include "tools.h"
#include "common.h"

/**
 * 使用OpenSSL生成密码学安全的随机数
 * @param num_bytes 指定要生成的随机数长度
 */
uint8_t* generate_secure_random_number(int num_bytes) {
    static uint8_t inited = 0;
    if (!inited) {
        RAND_poll();
        inited = 1;
    }
    uint8_t *buf = (uint8_t*)malloc(num_bytes);
    if (RAND_bytes(buf, num_bytes) != 1) {
        return NULL;
    }

    return buf;
}

/**
 * 使用三个随机数生成对称密钥
 */
uint8_t* generateSecret(uint8_t *A, uint8_t *B, uint8_t *C) {
    uint8_t *key = (uint8_t*)malloc(KEY_LENGTH);
    for (int i = 0; i < KEY_LENGTH; i++) {
        key[i] = A[i] ^ B[i] ^ C[i]; // 三个随机数进行异或操作，生成密钥
    }
    return key;
}


int encryptAndSign(uint8_t **out, uint32_t *out_len, uint8_t *in, uint32_t in_len, 
                    EC_KEY *encrypt_key, EC_KEY *sign_key, uint32_t uuid) {
     // 加密返回数据 sig_len(4) | sig | enc_data
    uint8_t *enc_data = NULL, *sig = NULL;
    uint8_t id[16] = {0};
    uint8_t digest[32] = {0};
    size_t enc_len = 0, sig_len = 0;

    memcpy(in, &uuid, sizeof(uint32_t));
    if (GM_SM2_encrypt(&enc_data, &enc_len, in, in_len, encrypt_key) != SUCCESS) {
        return FAILURE;
    }

    if (GM_SM3_digest(digest, enc_data, enc_len) != SUCCESS) {
        free(enc_data);
        return FAILURE;
    }

    if (GM_SM2_sign(&sig, &sig_len, digest, 32, id, sign_key) != SUCCESS) {
        free(enc_data);
        free(sig);
        return FAILURE;
    }
    *out_len = enc_len + sig_len + sizeof(uint32_t);
    *out = (uint8_t*)malloc(*out_len);
    memcpy(*out, &sig_len, sizeof(uint32_t));
    memcpy(*out + sizeof(uint32_t), sig, sig_len);
    memcpy(*out + sizeof(uint32_t) + sig_len, enc_data, enc_len);
    free(enc_data);
    free(sig);

    return SUCCESS;
}

int verifyAndDecrypt(uint8_t **out, uint32_t *out_len, uint8_t *in, uint32_t in_len, 
                EC_KEY *decrypt_key, EC_KEY *verify_key, uint32_t uuid) {
 // sig_len | sig | data
    uint32_t sig_len = 0, data_size = 0;
    // 验签
    uint8_t id[16] = {0};
    uint8_t digest[32] = {0};
    memcpy(&sig_len, in, sizeof(uint32_t));
    int offset = sizeof(uint32_t) + sig_len;
    uint32_t data_len = in_len - sig_len - sizeof(uint32_t);

    if (GM_SM3_digest(digest, in + offset, data_len) != SUCCESS) {
        printf("digest\n");
        return FAILURE;
    }

    if (GM_SM2_verify(in + sizeof(uint32_t), sig_len, digest, 32, id, verify_key) != SUCCESS) {
        printf("GM_SM2_verify\n");
        return FAILURE;
    }

    if (GM_SM2_decrypt(out, (size_t*)out_len, in + offset, data_len, decrypt_key) != SUCCESS) {
        printf("GM_SM2_decrypt\n");
        return FAILURE;
    }

    // 比对uid是否解密成功
    uint32_t tmp;
    memcpy(&tmp, *out, sizeof(uint32_t));
    if (tmp != uuid) {
        free(*out);
        *out = NULL;
        return FAILURE;
    }

    return SUCCESS;
}

int encryptAndSign_SM4(uint8_t **out, uint32_t *out_len, uint8_t *in, uint32_t in_len, 
                uint8_t *encrypt_key, EC_KEY *sign_key, uint32_t uuid) {
    // 加密返回数据 sig_len(4) | sig | enc_data
    uint8_t *enc_data = NULL, *sig = NULL;
    uint8_t iv[16]  = {0}, id[16]  = {0}, digest[32] = {0};
    size_t enc_len = 0, sig_len = 0;

    memcpy(in, &uuid, sizeof(uint32_t));
    enc_data = (uint8_t*)malloc(in_len + 16);
    if (GM_SM4_encrypt(enc_data, (int*)&enc_len, in, in_len, encrypt_key, iv) != SUCCESS) {
        free(enc_data);
        return FAILURE;
    }

    if (GM_SM3_digest(digest, enc_data, enc_len) != SUCCESS) {
        free(enc_data);
        return FAILURE;
    }

    if (GM_SM2_sign(&sig, &sig_len, digest, 32, id, sign_key) != SUCCESS) {
        free(enc_data);
        free(sig);
        return FAILURE;
    }
    *out_len = enc_len + sig_len + sizeof(uint32_t);
    *out = (uint8_t*)malloc(*out_len);
    memcpy(*out, &sig_len, sizeof(uint32_t));
    memcpy(*out + sizeof(uint32_t), sig, sig_len);
    memcpy(*out + sizeof(uint32_t) + sig_len, enc_data, enc_len);
    free(enc_data);
    free(sig);

    return SUCCESS;
}

int verifyAndDecrypt_SM4(uint8_t **out, uint32_t *out_len, uint8_t *in, uint32_t in_len, 
                uint8_t *decrypt_key, EC_KEY *verify_key, uint32_t uuid) {
    // data_size | sig_len | sig | data
    uint32_t sig_len = 0, data_size = 0;
    // 验签
    uint8_t id[16] = {0}, iv[16] = {0};
    uint8_t digest[32] = {0};
    memcpy(&sig_len, in, sizeof(uint32_t));
    int offset = sizeof(uint32_t) + sig_len;
    uint32_t data_len = in_len - sig_len - sizeof(uint32_t);


    if (GM_SM3_digest(digest, in + offset, data_len) != SUCCESS) {
        return FAILURE;
    }

    if (GM_SM2_verify(in + sizeof(uint32_t), sig_len, digest, 32, id, verify_key) != SUCCESS) {
        return FAILURE;
    }
    *out = (uint8_t*)malloc(data_len + 16);
    if (GM_SM4_decrypt(*out, out_len, in + offset, data_len, decrypt_key, iv) != SUCCESS) {
        free(*out);
        *out = NULL;
        return FAILURE;
    }

    // 比对uid是否解密成功
    uint32_t tmp;
    memcpy(&tmp, *out, sizeof(uint32_t));
    if (tmp != uuid) {
        free(*out);
        *out = NULL;
        return FAILURE;
    }

    return SUCCESS;
}

uint8_t* resolveHexStr(char *str) {
    uint8_t *res = (uint8_t*)malloc(strlen(str)/2);
    for (int i = 0; i < strlen(str); i += 2) {
        uint8_t left = 0, right = 0;
        if (str[i] >= '0' && str[i] <= '9') {
            left = str[i] - '0';
        } else {
            left = str[i] - 'a' + 10;
        }
        if (str[i + 1] >= '0' && str[i + 1] <= '9') {
            right = str[i + 1] - '0';
        } else {
            right = str[i + 1] - 'a' + 10;
        }
        res[i/2] = right | (left << 4);
    }

    return res;
}

void printHex(uint8_t *hex, uint32_t len) {
    for(int i = 0; i < len; ++i) {
        printf("%02x", hex[i]);
    }
    printf("\n");
}